#include "kernel_operator.h"
using namespace AscendC;

// template <typename T>
class KernelBroadCast
{
    const uint32_t BUFFER_NUM = 2;
    using T = float;

public:
    __aicore__ inline KernelBroadCast() {}
    __aicore__ inline void Init(TPipe *pipe, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                int32_t y_dimensional,
                                int32_t *y_ndarray, int32_t *x1_ndarray, int32_t *x2_ndarray,
                                int32_t *y_sumndarray, int32_t *x1_sumndarray, int32_t *x2_sumndarray)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->y_dimensional = y_dimensional;

        this->y_ndarray = y_ndarray;
        this->x1_ndarray = x1_ndarray;
        this->x2_ndarray = x2_ndarray;
        this->y_sumndarray = y_sumndarray;
        this->x1_sumndarray = x1_sumndarray;
        this->x2_sumndarray = x2_sumndarray;

        x1Gm.SetGlobalBuffer((__gm__ T *)x1, 1);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2, 1);
        yGm.SetGlobalBuffer((__gm__ T *)y, 1);

        pipe->InitBuffer(inQueueX, BUFFER_NUM, sizeof(float));
    }
    __aicore__ inline void Process()
    {
        int dim = this->y_dimensional;
        LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();

        for (int j = 0; j < this->y_sumndarray[dim]; j++)
        {
            int x1_start = 0, x2_start = 0;
            for (int k = 0; k < dim; k++)
            {
                if (this->x1_ndarray[k] != 1)
                {
                    x1_start += this->x1_sumndarray[k] * (j / this->y_sumndarray[k] % this->y_ndarray[k]);
                }
                if (this->x2_ndarray[k] != 1)
                {
                    x2_start += this->x2_sumndarray[k] * (j / this->y_sumndarray[k] % this->y_ndarray[k]);
                }
            }
            T x1 = x1Gm.GetValue(x1_start);
            T x2 = x2Gm.GetValue(x2_start);
            if constexpr (std::is_same_v<T, bfloat16_t>){
                xLocal.SetValue(0, ToFloat(x1));
                Ln(xLocal, xLocal, 1);
                Muls(xLocal, xLocal, ToFloat(x2), 1);
                Exp(xLocal, xLocal, 1);
                yGm.SetValue(j, ToBfloat16(xLocal.GetValue(0)));
            }else{
                xLocal.SetValue(0, static_cast<float>(x1)); 
                Ln(xLocal, xLocal, 1);
                Muls(xLocal, xLocal, static_cast<float>(x2), 1);
                Exp(xLocal, xLocal, 1);
                yGm.SetValue(j, static_cast<T>(xLocal.GetValue(0)));
            }
        }

        // int num = this->y_sumndarray[dim] / 512;
        // int res = this->y_sumndarray[dim] % 512;
        // for (int j = 0; j < num; j++)
        // {
        //     for(int i=0; i<512; i++){
        //         int x1_start = 0, x2_start = 0;
        //         for (int k = 0; k < dim; k++)
        //         {
        //             if (this->x1_ndarray[k] != 1)
        //             {
        //                 x1_start += this->x1_sumndarray[k] * ((j*512+i) / this->y_sumndarray[k] % this->y_ndarray[k]);
        //             }
        //             if (this->x2_ndarray[k] != 1)
        //             {
        //                 x2_start += this->x2_sumndarray[k] * ((j*512+i) / this->y_sumndarray[k] % this->y_ndarray[k]);
        //             }
        //         }
        //         T x1 = x1Gm.GetValue(x1_start);
        //         T x2 = x2Gm.GetValue(x2_start);
        //         if constexpr (std::is_same_v<T, bfloat16_t>){
        //             xLocal.SetValue(i, ToFloat(x1));
        //             x2Local.SetValue(i, ToFloat(x2));
                    
        //         }else{
        //             xLocal.SetValue(i, static_cast<float>(x1));
        //             x2Local.SetValue(i, static_cast<float>(x2));
                    
        //         }
        //     }
        //     Ln(xLocal, xLocal, 1);
        //     Mul(xLocal, xLocal, x2Local, 512);
        //     Exp(xLocal, xLocal, 512);
            
        //     yGm.SetValue(j, ToBfloat16(xLocal.GetValue(0)));
        // }

        // for (int j = 0; j < res; j++)
        // {
        //     int x1_start = 0, x2_start = 0;
        //     for (int k = 0; k < dim; k++)
        //     {
        //         if (this->x1_ndarray[k] != 1)
        //         {
        //             x1_start += this->x1_sumndarray[k] * ((num*512+j) / this->y_sumndarray[k] % this->y_ndarray[k]);
        //         }
        //         if (this->x2_ndarray[k] != 1)
        //         {
        //             x2_start += this->x2_sumndarray[k] * ((num*512+j) / this->y_sumndarray[k] % this->y_ndarray[k]);
        //         }
        //     }
        //     T x1 = x1Gm.GetValue(x1_start);
        //     T x2 = x2Gm.GetValue(x2_start);
        //     if constexpr (std::is_same_v<T, bfloat16_t>){
        //         xLocal.SetValue(0, ToFloat(x1));
        //         Ln(xLocal, xLocal, 1);
        //         Muls(xLocal, xLocal, ToFloat(x2), 1);
        //         Exp(xLocal, xLocal, 1);
        //         yGm.SetValue((num*512+j), ToBfloat16(xLocal.GetValue(0)));
        //     }else{
        //         xLocal.SetValue(0, static_cast<float>(x1)); 
        //         Ln(xLocal, xLocal, 1);
        //         Muls(xLocal, xLocal, static_cast<float>(x2), 1);
        //         Exp(xLocal, xLocal, 1);
        //         yGm.SetValue((num*512+j), static_cast<T>(xLocal.GetValue(0)));
        //     }
        // }
        inQueueX.FreeTensor(xLocal);
    }

private:
    GlobalTensor<T> x1Gm;
    GlobalTensor<T> x2Gm;
    GlobalTensor<T> yGm;

    int32_t y_dimensional;
    int32_t *y_ndarray;
    int32_t *x1_ndarray;
    int32_t *x2_ndarray;

    int32_t *y_sumndarray;
    int32_t *x1_sumndarray;
    int32_t *x2_sumndarray;

    TQue<QuePosition::VECIN, 1> inQueueX;
};
